# Created by 黄景涛
# DATE 2024/7/24

""" 重置组织下所有用户的密码，并写入yaml文件 """

import concurrent
from concurrent.futures import ThreadPoolExecutor

import requests

from user_manage_login import login_user_manage
from org_settings import fk
import math
import yaml

from config import settings

# 请求头配置
headers = {
    "User-Agent": fk.user_agent(),
    "Cookie": "bk_token=F_TI_c4CyRJQouc4h331WqEUhFOBMp5pZsVtDZnkLJw; blueking_language=zh-cn; bklogin_csrftoken=2TBTxEYRko4o4pCSirLTH0flJR0nWwOroc4zLcRTcKGkWkZV4mrlsjMtYBDmUNVN; cc-portal_saas_csrftoken=dNcTVUyGDDHdbSqfug1yu2XUfv1IK46htN9VpAnA9IVyzhBtYih7UJi3YNnUHUMr; cc-portal_saas_sessionid=tmg7y4qmi5xd8zr2zbf62j509hqh694k; ops-digital_saas_csrftoken=An32nRlMuHOEg0y8s18UW8vTcSLveqJVu8gvWUR28Z5igJYHHuaYNdA5tGvm4ywR; ops-digital_saas_sessionid=sszzry1nazrso0wmwt3ps8mu3is56io2; bk_csrftoken=DuvHtIUmmwfJEnsEATpFW1j1efgkvcmatYFiFXIC6w7eahukf4vR75DTLT9USv2Q; cc-portal-dev_saas_csrftoken=4Fz1xNVHEeEpgFGE94lDuQXHMcvmCHzPYuFiWNGHRHtpjFcQHAwx2dTQriNoqT1D; bkuser_csrftoken=j48bxLmbgNWHa31Sm57iwGAIiGAhIwoFxlSjrUsr3LpHpZgkaO05nJLPaojcOQC9; sessionid=13ewrix57kaqy9pbkc2b5z6jsih05kjx; o_f=%7B%22lock%22%3Afalse%2C%22msg%22%3A%7B%7D%7D; r_s=%7B%22lock%22%3Afalse%2C%22msg%22%3A%7B%7D%7D; r_f=%7B%22lock%22%3Atrue%2C%22msg%22%3A%7B%22workbench%22%3A2%7D%7D; o_s=%7B%22lock%22%3Atrue%2C%22msg%22%3A%7B%22workbench%22%3A%7B%222025-3-24%22%3A15494%7D%7D%7D",
    "X-Csrftoken": "j48bxLmbgNWHa31Sm57iwGAIiGAhIwoFxlSjrUsr3LpHpZgkaO05nJLPaojcOQC9"
}

# 获取配置
cfg = settings.cfg


def update_headers():
    base_url = cfg.get('base_url')
    user_home_url = cfg.get('user_manager').get('user_base_url') + settings.cfg.get('user_manager').get(
        'home_route')
    username = cfg.get('superman').get('username')
    password = cfg.get('superman').get('password')
    cookies, token = login_user_manage(base_url, user_home_url, username, password)
    headers['Cookie'] = cookies
    headers['X-Csrftoken'] = token


# update_headers()

# full_url
full_reset_password = cfg.get('user_manager').get('user_base_url') + cfg.get('user_manager').get('reset_password_route')
full_users_list = cfg.get('user_manager').get('user_base_url') + cfg.get('user_manager').get('search_users_list_route')

# 全局的session对象
session = requests.session()


def search_users_from_org(org_id):
    full_users_list_url = full_users_list + f'/{org_id}/profiles/'

    users = []
    page = 1
    page_size = 100
    params = {
        "page_size": page_size,
        "page": page,
        "recursive": True
    }

    try:
        resp = session.get(full_users_list_url, params=params, headers=headers)
        all_counts = resp.json().get('data').get('count')
        all_pages = math.ceil(all_counts / page_size)
        for p in range(1, all_pages + 1):
            param = {
                "page_size": page_size,
                "page": p,
                "recursive": True
            }
            resp = session.get(full_users_list_url, params=param, headers=headers)
            users_list = resp.json().get('data').get('data')
            for user in users_list:
                users.append((user['id'], user['display_name'], user['username']))
    except:
        raise Exception(f'查询组织{org_id}的用户列表失败')
    return users


# 重置当前用户的密码
def reset_password(user_info, new_password, headers):
    user_id, display_name, _ = user_info
    full_reset_password_url = full_reset_password + f'{user_id}/'
    data = {'password': new_password}
    try:
        resp = session.patch(full_reset_password_url, json=data, headers=headers)
        result = resp.json().get('result')
    except:
        raise Exception('重置密码的接口访问失败！')

    if result:
        print(f'用户【{display_name}】密码重置成功！')
    else:
        raise Exception(f'用户【{display_name}】密码重置失败！')


def run(org_id, new_passwd):
    users = search_users_from_org(org_id)
    data = {}
    for user in users:
        try:
            reset_password(user, new_passwd, headers)
        except Exception as e:
            print(e)
        else:
            data.update([(user[-1], new_passwd)])

    with open('user_files/users.yaml', mode='a+', encoding='utf8') as f:
        yaml.dump(data, f)


def concurrent_run(org_id, new_passwd):
    users = search_users_from_org(org_id)
    data = {}
    # 使用线程池来并发执行重置密码的操作
    with ThreadPoolExecutor(max_workers=10) as executor:  # 可以根据需要调整max_workers的值
        futures = {executor.submit(reset_password, user, new_passwd, headers): user for user in users}

        for future in concurrent.futures.as_completed(futures):
            user = futures[future]
            try:
                future.result()  # 获取线程执行的结果
                data.update([(user[-1], new_passwd)])
            except Exception as e:
                print(e)

    with open('user_files/users.yaml', mode='a+', encoding='utf8') as f:
        yaml.dump(data, f)


if __name__ == '__main__':
    # 272
    org_id = "9177"
    new_passwd = "Bk@123456"
    concurrent_run(org_id, new_passwd)
